import os
import time
import shutil
from osgeo import gdal
from common.io_utils.downloader import download_via_wget
from common.io_utils.dirs import mkdir_p, chmod_path
from common.logger import logger
from common.geoimage.gdal_sys import SceneMeta, ValueInterpretationMeta
from common.geoimage.raster_dataset import RasterDataset
from catalog.config import DIR_STAC_BASE, DIR_DATA_BASE
from peewee import Query
from common.typing_utils import numeric
from typing import Optional, Tuple

ROOT_PATH = "paii-meta://catalog.json"

NAIP_BASE_MSFT_URL = "https://naipeuwest.blob.core.windows.net/naip/v002"
NAIP_BASE_NAS_DIR = "/NAS/DATA/USDA/NAIP"
NAIP_BASE_CACHE_DIR = "/NAS6/Members/skywalker/shared_storage"
NAIP_DEFAULT_CACHE_DIR = f"{NAIP_BASE_CACHE_DIR}/{os.getuid()}/.skywalker/cache"

gdal.UseExceptions()


def _transfer_from_msft_naip(url_src_item, dir_item, thumbnail: bool = False):
    """Transfer NAIP collection from Microsoft Planetary Computer

    `/NAS` data storage for `NAIP`
    |- /NAS/DATA/USDA/NAIP/.../{Item name}
    |-- R.tif
    |-- G.tif
    |-- B.tif
    |-- N.tif
    |-- thumbnail.200.jpg

    Disable downloading "thumbnail" by default.
    It's because most thumbnails are not found in Microsoft Planetary Computer.
    """
    mkdir_p(dir_item)
    item_name = os.path.basename(dir_item)

    try:
        t0 = time.time()
        # Download image, each band saved as a single file
        with chmod_path(dir_item, mode=0o777, mode_temp=0o777, child_mode=0o777):
            ds = gdal.Open("/vsicurl/" + url_src_item)
            band_meta = SceneMeta(
                n_cols=ds.RasterXSize,
                n_rows=ds.RasterYSize,
                n_bands=1,
                prj_wkt=ds.GetProjection(),
                geotrans=ds.GetGeoTransform(),
                value_interpretations=[ValueInterpretationMeta(1, 0, 0)],
            )
            # Download RGBN bands separately
            bands_map = {"R": 0, "G": 1, "B": 2, "N": 3}
            for band, idx in bands_map.items():
                RasterDataset.from_ndarray(
                    ds.GetRasterBand(idx + 1).ReadAsArray(), band_meta
                ).to_geotiff(os.path.join(dir_item, f"{band}.tif"))

            logger.info(
                f"Done. Downloaded item `{item_name}` in {(time.time() - t0) / 60:.2f}m."
            )
            if thumbnail:
                # Download thumbnail
                ret_thumbnail = download_via_wget(
                    url_src_item.replace(".tif", ".200.jpg"),
                    dest_dir=dir_item,
                    dest_file_name="thumbnail.200.jpg",
                    verbose_level=0,
                )
                if ret_thumbnail.returncode == 0:
                    logger.info(f"Downloaded thumbnail `{item_name}/thumbnail.200.jpg`")
                else:
                    logger.warning(
                        f"Failed to download thumbnail `{item_name}/thumbnail.200.jpg`"
                    )

    except Exception as e:
        shutil.rmtree(dir_item)
        logger.warning(
            f"=== Failed to download `{item_name}` from `{url_src_item}` due to err `{e}` ==="
        )
    finally:
        ds = None


def interpret_protocol(
    url: str, dir_base_stac=DIR_STAC_BASE, dir_base_data=DIR_DATA_BASE
) -> str:
    """Interpret custom defined protocol used in `url`

    Note
    ----
    - map `paii-meta://` to `DIR_STAC_BASE`, and vice-versa
    - map `paii-data://` to `DIR_DATA_BASE`, and vice-versa
    - you could change the directory location by overwrite via `dir_base_stac` and `dir_base_data`
    - if the given url is not started with a known custom protocol or known directory path, return the url unchanged
    """

    if dir_base_stac is None:
        dir_base_stac = DIR_STAC_BASE
    if dir_base_data is None:
        dir_base_data = DIR_DATA_BASE

    if not dir_base_stac.endswith("/"):
        dir_base_stac += "/"
    if not dir_base_data.endswith("/"):
        dir_base_data += "/"

    if url.startswith("paii-meta://"):
        url = url.replace("paii-meta://", dir_base_stac)
    elif url.startswith("paii-data://"):
        url = url.replace("paii-data://", dir_base_data)
    elif url.startswith(dir_base_stac):
        url = url.replace(dir_base_stac, "paii-meta://")
    elif url.startswith(dir_base_data):
        url = url.replace(dir_base_data, "paii-data://")
    elif url.startswith(NAIP_BASE_MSFT_URL):
        # url e.g.: https://naipeuwest.blob.core.windows.net/naip/v002
        # /fl/2019/fl_60cm_2019/28080/m_2808060_sw_17_060_20191215/R.tif
        url_src_item = url[: -len("/R.tif")] + ".tif"
        url = url.replace(NAIP_BASE_MSFT_URL, NAIP_BASE_NAS_DIR)
        if not os.path.exists(url):
            # Downloaded data to `DEFAULT_CACHE_DIR`. nas account will move data to /NAS periodically.
            url = NAIP_DEFAULT_CACHE_DIR + url
            if not os.path.exists(url):
                dir_item = os.path.dirname(url)
                _transfer_from_msft_naip(url_src_item, dir_item)

    return url


def helper_query_range(
    query: Query,
    n_conditions: int,
    q_object,
    q_range: Tuple[numeric, numeric],
    except_value: Optional[int] = None,
) -> tuple[Query, int]:
    """A helper for SQL range query

    Parameters
    ----------
    query : Query
        The object of Query
    n_conditions : int
        The number of query conditions
    q_object : _type_
        The attribute to be queried
    q_range : Tuple[numeric, numeric]
        The range of query condistions, (lower, upper)
    except_value : int, optional
        The exceptional value of the query attribute, by default None
        The initial indexing table contains -1 value that we set except_value=-1 to handle.

    Returns
    -------
    tuple[Query, int]
        Return the query object and number of query conditions

    Exceptions
        Raise ValueError if upper bound is less than the lower bound.
    """
    if (not isinstance(q_range, tuple) and not isinstance(q_range, list)) or len(
        q_range
    ) != 2:
        msg = "[Error] Query conditions should be in tuple or list with length of 2, (low, high)."
        raise Exception(msg + f" Receive the input {q_range}")

    lower, upper = q_range
    if lower is not None and upper is not None and upper < lower:
        msg = f"[Error] The query upper bound ({upper}) should be larger than the lower bound ({lower})."
        raise Exception(msg)

    if lower is not None:
        if except_value is not None:
            query = query.where((q_object >= lower) | (q_object == except_value))
        else:
            query = query.where(q_object >= lower)
        n_conditions += 1

    if upper is not None:
        if except_value is not None:
            query = query.where((q_object <= upper) | (q_object == except_value))
        else:
            query = query.where(q_object <= upper)
        n_conditions += 1

    return query, n_conditions
